package main

import (
	"bufio"
	"fmt"
	"go2/log"
	pool2 "go2/pool"
	"io"
	"net"
	"sync"
	"time"
)

func handleSOCKS5(conn io.Writer, reader *bufio.Reader) {
	// Use single small buffer for initial protocol exchange
	protocolBuf := pool2.Malloc(256)
	defer pool2.Free(protocolBuf)

	// Read version and number of methods
	if _, err := io.ReadFull(reader, protocolBuf[:2]); err != nil {
		return
	}

	version, nmethods := protocolBuf[0], protocolBuf[1]
	if version != 0x05 {
		return
	}

	// Read authentication methods
	if _, err := io.ReadFull(reader, protocolBuf[2:2+nmethods]); err != nil {
		return
	}

	// Send no authentication response
	if _, err := conn.Write([]byte{0x05, 0x00}); err != nil {
		return
	}

	// Read connection request
	if _, err := io.ReadFull(reader, protocolBuf[:4]); err != nil {
		return
	}

	version, cmd, atyp := protocolBuf[0], protocolBuf[1], protocolBuf[3]
	if version != 0x05 || cmd != 0x01 {
		conn.Write([]byte{0x05, 0x07, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00})
		return
	}

	var addr string
	switch atyp {
	case 0x01: // IPv4
		if _, err := io.ReadFull(reader, protocolBuf[4:8]); err != nil {
			return
		}
		addr = net.IP(protocolBuf[4:8]).String()
	case 0x03: // Domain name
		if _, err := io.ReadFull(reader, protocolBuf[4:5]); err != nil {
			return
		}
		domainLen := int(protocolBuf[4])
		if domainLen > 255 { // Prevent buffer overflow
			return
		}
		if _, err := io.ReadFull(reader, protocolBuf[5:5+domainLen]); err != nil {
			return
		}
		addr = string(protocolBuf[5:5+domainLen])
	case 0x04: // IPv6
		if _, err := io.ReadFull(reader, protocolBuf[4:20]); err != nil {
			return
		}
		addr = net.IP(protocolBuf[4:20]).String()
	default:
		conn.Write([]byte{0x05, 0x08, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00})
		return
	}

	// Read port
	portBuf := protocolBuf[:2]
	if _, err := io.ReadFull(reader, portBuf); err != nil {
		return
	}
	portNum := (int(portBuf[0]) << 8) | int(portBuf[1])

	remote_addr := fmt.Sprintf("%s:%d", addr, portNum)

	log.Debugf("handleSOCKS5: %s", remote_addr)

	// Set up dial with timeout for better resource management
	dialer := net.Dialer{Timeout: 30 * time.Second}
	target, err := dialer.Dial("tcp", remote_addr)
	if err != nil {
		conn.Write([]byte{0x05, 0x01, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00})
		return
	}
	defer target.Close()

	// Send success response
	if _, err := conn.Write([]byte{0x05, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00}); err != nil {
		return
	}

	// Use larger buffers for data transfer
	buf1 := pool2.Malloc(64 * 1024) // Increased to 64KB for better throughput
	defer pool2.Free(buf1)
	buf2 := pool2.Malloc(64 * 1024) // Increased to 64KB for better throughput
	defer pool2.Free(buf2)

	// Use WaitGroup for proper cleanup
	var wg sync.WaitGroup
	wg.Add(1)

	go func() {
		defer wg.Done()
		_, _ = io.CopyBuffer(target, reader, buf1)
	}()

	_, _ = io.CopyBuffer(conn, target, buf2)
	wg.Wait() // Ensure the copy goroutine completes
}
